import warnings
warnings.filterwarnings('ignore')
import util_for_output_zh

import os
import pandas as pd
import numpy as np
# import matplotlib.pyplot as plt
from util_set_zh_matplot import plt
import seaborn as sns
from sklearn.preprocessing import LabelEncoder, StandardScaler
from sklearn.model_selection import train_test_split, StratifiedKFold, cross_val_score
from sklearn.metrics import classification_report, confusion_matrix, roc_auc_score
from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
import xgboost as xgb
from imblearn.over_sampling import SMOTE
from imblearn.pipeline import Pipeline
import joblib

# 设置显示选项
pd.set_option('display.max_columns', None)
pd.set_option('display.width', None)

# 创建输出目录
out_dir = 'ques4_female_fetus_results'
os.makedirs(out_dir, exist_ok=True)

'''
第四问
问题：女胎数据在 染色体的非整倍体 每个分类上 1 数据量少 2 数据量不均衡
针对数据量少 准备加入男胎数据
            出现问题 男女胎数据一起训练效果差
            解决方案 1 相关性分析 2 使用男胎数据训练预测女胎 3 女胎数据训练预测男胎
            发现 X 与性别相关度太高 且 X的相关信息 影响了 从男胎数据中 学习 女胎知识
            因此，将X相关信息 排除在 输入的特征以外
    问题 男胎数据并不是最终要求解的问题，而是通过男胎数据观测女胎数据的一部分特征
    解决 使用迁移学习的思路，将男胎数据（粗糙）全部+女胎数据（准确）一部分 作为训练集
            使得模型在男胎数据上初步学习 “相关性” 知识，而女胎数据则 微调 这部分知识
针对数据量不均衡
    使用SMOTE对数据量多的做欠采样 数据量少的做过采样 使得 样本量均衡
            
最终 使用这种方式训练的模型，在女胎 测试集上获得较好的效果

'''


def load_and_preprocess_data(file_path):
    """加载并预处理数据，分离男女样本并筛选类别"""
    # 加载数据
    data = pd.read_csv(file_path)
    
    # 移除性别相关特征（避免男性数据干扰女胎预测）
    gender_related_features = ['X染色体的Z值', 'X染色体浓度', 'Y染色体的Z值', 'Y染色体浓度']
    data = data.drop(columns=[f for f in gender_related_features if f in data.columns])
    
    # 分离男女数据
    female_data = data[data['性别'] == '女'].copy()
    male_data = data[data['性别'] == '男'].copy()
    
    # 获取女胎和男胎中类别数量少于10的类别
    female_counts = female_data['染色体的非整倍体'].value_counts()
    male_counts = male_data['染色体的非整倍体'].value_counts()
    
    # 找出两类数据中都少于10的类别
    rare_female_categories = female_counts[female_counts < 10].index
    rare_male_categories = male_counts[male_counts < 10].index
    rare_categories = set(rare_female_categories).union(set(rare_male_categories))
    
    # 筛选掉这些稀有类别
    female_data = female_data[~female_data['染色体的非整倍体'].isin(rare_categories)]
    male_data = male_data[~male_data['染色体的非整倍体'].isin(rare_categories)]
    
    print(f"加载女胎数据 {len(female_data)} 条，男胎数据 {len(male_data)} 条")
    print(f"移除的稀有类别: {rare_categories}")
    
    # 处理缺失值
    numeric_cols = female_data.select_dtypes(include=['float64', 'int64']).columns
    female_data[numeric_cols] = female_data[numeric_cols].fillna(female_data[numeric_cols].median())
    male_data[numeric_cols] = male_data[numeric_cols].fillna(male_data[numeric_cols].median())

    # 查看女胎类别分布
    print("\n女胎类别分布:")
    print(female_data['染色体的非整倍体'].value_counts())
    
    # 可视化女胎类别分布
    plt.figure(figsize=(10, 6))
    ax = sns.countplot(x='染色体的非整倍体', data=female_data)
    plt.title('女胎染色体非整倍体类别分布')
    plt.xticks(rotation=45)

    # 在每个柱子上添加数值
    for p in ax.patches:
        ax.annotate(f'{int(p.get_height())}', 
                    (p.get_x() + p.get_width() / 2., p.get_height()), 
                    ha='center', va='center', 
                    xytext=(0, 5), 
                    textcoords='offset points')

    plt.tight_layout()
    plt.savefig(os.path.join(out_dir, '女胎类别分布.png'))
    plt.close()
    
    return female_data, male_data


def select_features(data):
    """选择与目标染色体相关的特征"""
    # 核心特征：目标染色体相关特征（移除性别相关特征）
    features = [
        '13号染色体的Z值', '18号染色体的Z值', '21号染色体的Z值',
        '13号染色体的GC含量', '18号染色体的GC含量', '21号染色体的GC含量',
        '原始读段数', '在参考基因组上比对的比例', '重复读段的比例', 
        '唯一比对的读段数', 'GC含量'
    ]
    
    # 检查特征存在性
    existing_features = [f for f in features if f in data.columns]
    
    # 特征缺失提示
    missing_features = set(features) - set(existing_features)
    if missing_features:
        print(f"警告：以下特征不存在于数据中，已自动忽略: {missing_features}")
    
    return existing_features

def prepare_training_data(female_data, male_data, features, test_size=0.3):
    """准备训练集和测试集：测试集全为女性，训练集包含部分女性和大量男性"""
    # 确保类别一致
    female_classes = set(female_data['染色体的非整倍体'].unique())
    male_data = male_data[male_data['染色体的非整倍体'].isin(female_classes)].copy()
    
    # 划分女胎数据为训练和测试（测试集全为女性）
    female_train, female_test = train_test_split(
        female_data, test_size=test_size, random_state=42, 
        stratify=female_data['染色体的非整倍体']
    )
    
    # 训练集：合并部分女性和大量男性数据（保持性别比例合理）
    # 根据女胎训练集大小按比例选取男性数据
    male_sample_size = min(len(male_data), int(len(female_train) * 1.5))  # 男性数据为女性1.5倍
    
    # 修正：使用train_test_split实现男性数据的分层抽样（兼容旧pandas版本）
    if len(male_data) > 0 and male_sample_size > 0:
        # 用train_test_split实现分层抽样
        _, male_train = train_test_split(
            male_data, 
            test_size=male_sample_size/len(male_data),  # 按比例抽样
            random_state=42,
            stratify=male_data['染色体的非整倍体']
        )
    else:
        male_train = pd.DataFrame(columns=male_data.columns)  # 空数据框
    
    # 合并训练集
    train_data = pd.concat([female_train, male_train], ignore_index=True)
    
    print(f"\n训练集组成: 女胎 {len(female_train)} 条, 男胎 {len(male_train)} 条, 共 {len(train_data)} 条")
    print(f"测试集组成: 女胎 {len(female_test)} 条 (全女性)")
    print("\n训练集类别分布:")
    print(train_data['染色体的非整倍体'].value_counts())
    
    # 准备特征和标签
    X_train = train_data[features]
    y_train = train_data['染色体的非整倍体']
    X_test = female_test[features]
    y_test = female_test['染色体的非整倍体']
    
    return X_train, y_train, X_test, y_test



def transfer_learning_model(X_pretrain, y_pretrain, X_finetune, y_finetune, X_test, y_test, class_names):
    """迁移学习模型：先在整体数据预训练，再用女胎数据微调"""
    # 数据标准化
    scaler = StandardScaler()
    X_pretrain_scaled = scaler.fit_transform(X_pretrain)
    X_finetune_scaled = scaler.transform(X_finetune)
    X_test_scaled = scaler.transform(X_test)
    
    # 1. 预训练模型（使用全部可用数据）
    pretrain_model = xgb.XGBClassifier(
        objective='multi:softmax',
        num_class=len(class_names),
        n_estimators=150,
        max_depth=5,
        learning_rate=0.1,
        reg_alpha=0.1,
        reg_lambda=0.1,
        random_state=42
    )
    pretrain_model.fit(X_pretrain_scaled, y_pretrain)
    
    # 2. 微调模型（使用女胎数据）
    # 降低学习率进行微调
    finetune_model = xgb.XGBClassifier(
        objective='multi:softmax',
        num_class=len(class_names),
        n_estimators=100,
        max_depth=5,
        learning_rate=0.01,  # 小学习率微调
        reg_alpha=0.1,
        reg_lambda=0.1,
        random_state=42
    )
    
    # 加载预训练权重
    # finetune_model.load_model(pretrain_model.save_raw())
    finetune_model._Booster = pretrain_model.get_booster()
    
    # 处理女胎训练集中的类别不平衡
    # 动态调整SMOTE参数
    unique_classes, class_counts = np.unique(y_finetune, return_counts=True)
    min_samples = np.min(class_counts)
    n_neighbors = min(5, min_samples - 1) if min_samples > 1 else 1
    
    if min_samples <= 1:
        print("警告：部分女胎类别样本极少，使用原始数据进行微调")
        X_finetune_resampled, y_finetune_resampled = X_finetune_scaled, y_finetune
    else:
        print("SMOTE 过采样前数据量:")
        print(f"X_finetune_scaled shape: {X_finetune_scaled.shape}")
        print(f"y_finetune shape: {y_finetune.shape}")
        # print(f"类别分布（y_finetune）:\n{y_finetune.value_counts()}")

        smote = SMOTE(random_state=42, k_neighbors=n_neighbors)
        X_finetune_resampled, y_finetune_resampled = smote.fit_resample(
            X_finetune_scaled, y_finetune
        )

        print("\nSMOTE 过采样后数据量:")
        print(f"X_finetune_resampled shape: {X_finetune_resampled.shape}")
        print(f"y_finteune_resampled shape: {y_finetune_resampled.shape}")
        # print(f"类别分布（y_finetune_resampled）:\n{y_finetune_resampled.value_counts()}")
    
    # 微调模型
    finetune_model.fit(
        X_finetune_resampled, 
        y_finetune_resampled,
        # xgb_model=finetune_model.save_raw()  # 基于预训练权重继续训练
        xgb_model=finetune_model.get_booster()
    )
    
    # 评估微调后的模型
    y_pred = finetune_model.predict(X_test_scaled)
    
    # 交叉验证评估
    cv = StratifiedKFold(n_splits=min(5, len(X_finetune_resampled)//5), shuffle=True, random_state=42)
    cv_scores = cross_val_score(
        finetune_model, 
        X_finetune_resampled, 
        y_finetune_resampled, 
        cv=cv, 
        scoring='f1_macro'
    )
    print(f"\n交叉验证F1分数: {cv_scores.mean():.4f} ± {cv_scores.std():.4f}")
    
    return finetune_model, scaler, y_test, y_pred


def evaluate_model(y_true, y_pred, class_names):
    """评估模型性能并可视化结果"""
    # 打印分类报告
    print("\n模型评估报告:")
    print(classification_report(y_true, y_pred, target_names=class_names))
    
    # 绘制混淆矩阵
    plt.figure(figsize=(10, 8))
    cm = confusion_matrix(y_true, y_pred)
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',
                xticklabels=class_names, yticklabels=class_names)
    plt.xlabel('预测类别')
    plt.ylabel('实际类别')
    plt.title('女胎异常判定混淆矩阵')
    plt.tight_layout()
    plt.savefig(os.path.join(out_dir, '混淆矩阵.png'))
    plt.close()


def analyze_feature_importance(model, features):
    """分析并可视化特征重要性"""
    importance = pd.DataFrame({
        '特征': features,
        '重要性': model.feature_importances_
    }).sort_values('重要性', ascending=False)
    
    print("\n特征重要性排序:")
    print(importance)
    
    # 可视化特征重要性
    plt.figure(figsize=(12, 6))
    sns.barplot(x='重要性', y='特征', data=importance)
    plt.title('特征重要性（女胎异常判定）')
    plt.tight_layout()
    plt.savefig(os.path.join(out_dir, '特征重要性.png'))
    plt.close()
    
    return importance


def main():
    # 1. 加载并预处理数据
    data_path = 'ques4_prepare_data/ques4_concat.csv'
    female_data, male_data = load_and_preprocess_data(data_path)
    
    # 2. 特征选择（移除性别相关特征）
    features = select_features(female_data)
    print(f"\n使用的特征: {features}")
    
    # 3. 准备训练和测试数据
    # 训练集：部分女性 + 大量男性；测试集：全女性
    X_train_mixed, y_train_mixed, X_test_female, y_test_female = prepare_training_data(
        female_data, male_data, features, test_size=0.3
    )
    
    # 4. 准备迁移学习的预训练和微调数据
    # 预训练数据：所有可用数据（男+女）
    X_pretrain = pd.concat([X_train_mixed, X_test_female], ignore_index=True)
    y_pretrain = pd.concat([y_train_mixed, y_test_female], ignore_index=True)
    
    # 微调数据：仅女胎训练数据
    X_finetune = X_train_mixed[X_train_mixed.index.isin(female_data.index)]
    y_finetune = y_train_mixed[X_train_mixed.index.isin(female_data.index)]
    
    # 标签编码
    le = LabelEncoder()
    le.fit(y_pretrain)  # 使用所有类别进行编码
    y_pretrain_encoded = le.transform(y_pretrain)
    y_finetune_encoded = le.transform(y_finetune)
    y_test_encoded = le.transform(y_test_female)
    class_names = le.classes_
    
    # 5. 构建迁移学习模型
    model, scaler, y_test, y_pred = transfer_learning_model(
        X_pretrain, y_pretrain_encoded,
        X_finetune, y_finetune_encoded,
        X_test_female, y_test_encoded,
        class_names
    )
    
    # 6. 评估模型性能
    evaluate_model(y_test, y_pred, class_names)
    
    # 7. 分析特征重要性
    analyze_feature_importance(model, features)
    
    # 8. 保存模型
    joblib.dump({
        'model': model,
        'scaler': scaler,
        'label_encoder': le,
        'features': features
    }, os.path.join(out_dir, 'female_fetus_transfer_model.pkl'))
    
    print("\n模型及分析结果已保存至 female_fetus_transfer_results 目录")


if __name__ == '__main__':
    main()